#pragma once

#include "scene.h"
#include "sampler.h"
#include "film.h"
#include "utils/singleton.h"
#include "utils/thread_pool.h"

#include <numeric>

namespace kp {

	struct IntegrateProgress {
		int current;
		int total;
		IntegrateProgress(int cur, int tol) :current(cur), total(tol) {}
		
		IntegrateProgress operator+(const IntegrateProgress& ip) {
			return { this->current + ip.current,this->total + ip.total };
		}
	};

	struct IntegrateState {
		size_t tasknum = 0;
		std::vector<IntegrateProgress> progress;

		IntegrateProgress progression() const {
			return std::accumulate(progress.begin(), progress.end(), IntegrateProgress(0, 0));
		}
	};

	class Integrator {
	public:
		Integrator() {}

		virtual void render(Ref<Scene> scene, Ref<Sampler> sampler, Ref<Film>film) = 0;
		
		IntegrateState state;

	protected:
		Ref<thread_pool> pool;

		Spectrum Li(Scene& scene, const Ray& ray, Ref<Sampler> sampler) {
			Intersection isct;
			if (scene.Intersect(ray, isct)) {
				isct.primitive->ComputeScatteringFunctions(&isct, TransportMode::Radiance, false);
				return Spectrum(isct.Ns.x + 1.0f, isct.Ns.y + 1.0f, isct.Ns.z + 1.0f) * 0.5f;
			}
			else {
				float t = 0.5f * (isct.Ng.y + 1.0f);
				return  (1.0f - t) * Spectrum(1.0f) + t * Spectrum(0.5f, 0.7f, 1.0f);
			}
		}

	};

	class Tracer : public Integrator {
	public:
		Tracer() {
			this->pool = singleton<thread_pool>::getInstance(1);
		}

		virtual void render(Ref<Scene> scene, Ref<Sampler> sampler, Ref<Film>film) override {
			is_started = true;
			auto width = film->size.w, height = film->size.h;
			state.tasknum = 1;
			state.progress.emplace_back(0, height * width * sampler->spp);
			pool->submit([=]() {
				auto fragColor = Spectrum(0.0f);
				if (!is_started)return;
				for (int h = 0; h < height; h++) {
					for (int w = 0; w < width; w++) {
						for (int s = 0; s < sampler->spp; s++) {
							if (!is_started)return;
							// Trace ray
							Ray ray;
							scene->camera->generateRay(sampler->getCameraSample({ w,h }), ray);
							fragColor += Li(*scene, ray, sampler);
							state.progress[0].current++;
						}
						fragColor /= (float)sampler->spp;
						// Write frag color
						film->frags[w + h * (size_t)width] = fragColor;
						fragColor = Spectrum(0.0f);
					}
				}
				});
		}

	private:
		bool is_started = false;
	};

	class MultiTracer : public Integrator {
	public:
		MultiTracer(size_t threadNumbers) {
			this->pool = singleton<thread_pool>::getInstance(threadNumbers);
		}

		virtual void render(Ref<Scene> scene, Ref<Sampler> sampler, Ref<Film>film) override {
			is_started = true;
			auto width = film->size.w, height = film->size.h;
			const size_t unit_height = height / pool->getThreadSize() / 4;
			size_t splitNum = height / unit_height;
			state.tasknum = height % unit_height ? splitNum + 1 : splitNum;
			size_t tmpHeight = height;
			for (int i = 0; i < state.tasknum; i++) {
				if (tmpHeight < unit_height) {
					state.progress.emplace_back(0, width * tmpHeight * sampler->spp);
				}
				else {
					state.progress.emplace_back(0, width * unit_height * sampler->spp);
				}
				tmpHeight -= unit_height;
			}
			int h_start = 0;
			for (int i = 0; i < state.tasknum; i++) {
				pool->submit([=](int start_h, int end_h) {
					auto fragColor = Spectrum(0.0f);
					if (!is_started)return;
					for (int h = start_h; h < height && h < end_h; h++) {
						for (int w = 0; w < width; w++) {
							for (int s = 0; s < sampler->spp; s++) {
								if (!is_started)return;
								// Get uv
								// Trace ray
								Ray ray;
								scene->camera->generateRay(sampler->getCameraSample({ w,h }), ray);
								fragColor += Li(*scene, ray, sampler);
								state.progress[i].current++;
							}
							fragColor /= (float)sampler->spp;
							// Write frag color
							film->frags[w + h * (size_t)width] = fragColor;
							fragColor = Spectrum(0.0f);
						}
					}
					}, h_start, h_start + unit_height);
				h_start += (int)unit_height;
			}

		}

	private:
		
		bool is_started = false;

	};

}
